import requests
import json

class ComfyUIClient:
    """ComfyUI客户端类，用于与ComfyUI服务器进行交互"""
    
    def __init__(self, server_url="http://localhost:8188"):
        """
        初始化ComfyUI客户端
        :param server_url: ComfyUI服务器地址，默认为http://localhost:8188
        """
        self.server_url = server_url
    
    def get_history(self, prompt_id):
        """
        获取指定prompt_id的历史记录
        :param prompt_id: 提示词ID
        :return: 历史记录的JSON数据
        """
        response = requests.get(f"{self.server_url}/history/{prompt_id}")
        return response.json()
    
    def post_prompt(self, workflow_json):
        """
        向服务器发送工作流
        :param workflow_json: 工作流JSON数据
        :return: 服务器响应的JSON数据
        """
        payload = {
            "prompt": workflow_json,
            # 包含 workflow 到 extra_pnginfo，避免服务端关于缺少 'workflow' 的报错/警告
            "extra_pnginfo": {"workflow": workflow_json},
        }
        response = requests.post(f"{self.server_url}/prompt", json=payload)
        return response.json()
    
    def get_image(self, filename, subfolder, folder_type):
        """
        获取图片数据
        :param filename: 文件名
        :param subfolder: 子文件夹名
        :param folder_type: 文件夹类型
        :return: 图片二进制数据
        """
        params = {"filename": filename, "subfolder": subfolder, "type": folder_type}
        response = requests.get(f"{self.server_url}/view", params=params)
        return response.content
    
    def load_workflow(self, file_path):
        """
        从文件加载工作流
        :param file_path: 工作流文件路径
        :return: 工作流JSON数据
        """
        with open(file_path, 'r', encoding='utf-8') as f:
            return json.load(f)